#include "RayTracingLibrary.h"
#include "../../Core/CoreGlobals.h"
#include "../../Utils/LogUtil.h"
#include "../../Templates/NewObject.h"
#include "../../Game/Process/RenderProcess.h"

RayTracingLibrary::Builder::Builder() noexcept = default;
RayTracingLibrary::Builder::~Builder() noexcept = default;
RayTracingLibrary::Builder::Builder(RayTracingLibrary::Builder const& rhs) noexcept = default;
RayTracingLibrary::Builder::Builder(RayTracingLibrary::Builder && rhs) noexcept = default;
RayTracingLibrary::Builder& RayTracingLibrary::Builder::operator=(RayTracingLibrary::Builder const& rhs) noexcept = default;
RayTracingLibrary::Builder& RayTracingLibrary::Builder::operator=(RayTracingLibrary::Builder && rhs) noexcept = default;

RayTracingLibrary::Builder& RayTracingLibrary::Builder::SetDescriptor(RefCountPtr<RHIDescriptor> descriptor) noexcept
{
	mImpl->descriptor = descriptor;
	return *this;
}

RayTracingLibrary::Builder& RayTracingLibrary::Builder::AddShaderStage(ShaderPtr<RHIShader> shader) noexcept
{
	// make sure shaderStage is add before shaderGroup
	assert(mImpl->shaderGroups.size() == 0);
	mImpl->shaderStages.push_back(shader);
	return *this;
}

RayTracingLibrary::Builder& RayTracingLibrary::Builder::PushShader(ShaderPtr<RHIShader> shader) noexcept
{
	switch (shader->GetShaderStage())
	{
	case VK_SHADER_STAGE_RAYGEN_BIT_KHR:
		PushRaygen(shader);
		break;
	case VK_SHADER_STAGE_MISS_BIT_KHR:
		PushMiss(shader);
		break;
	case VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR:
		PushClosestHit(shader);
		break;
	case VK_SHADER_STAGE_ANY_HIT_BIT_KHR:
		PushAnyHit(shader);
		break;
	case VK_SHADER_STAGE_INTERSECTION_BIT_KHR:
		PushIntersection(shader);
		break;
	case VK_SHADER_STAGE_CALLABLE_BIT_KHR:
		PushCallable(shader);
		break;
	default:
		assert(0);
		break;
	}
	return *this;
}

RayTracingLibrary::Builder& RayTracingLibrary::Builder::PushRaygen(ShaderPtr<RHIShader> shader) noexcept
{
	assert(shader->GetShaderStage() == VK_SHADER_STAGE_RAYGEN_BIT_KHR);
	assert(mImpl->tempCroup.generalShader == VK_SHADER_UNUSED_KHR);
	mImpl->tempCroup.generalShader = GetShaderIndex(shader);
	return *this;
}

RayTracingLibrary::Builder& RayTracingLibrary::Builder::PushMiss(ShaderPtr<RHIShader> shader) noexcept
{
	assert(shader->GetShaderStage() == VK_SHADER_STAGE_MISS_BIT_KHR);
	assert(mImpl->tempCroup.generalShader == VK_SHADER_UNUSED_KHR);
	mImpl->tempCroup.generalShader = GetShaderIndex(shader);
	return *this;
}

RayTracingLibrary::Builder& RayTracingLibrary::Builder::PushClosestHit(ShaderPtr<RHIShader> shader) noexcept
{
	assert(shader->GetShaderStage() == VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR);
	assert(mImpl->tempCroup.closestHitShader == VK_SHADER_UNUSED_KHR);
	mImpl->tempCroup.closestHitShader = GetShaderIndex(shader);
	return *this;
}

RayTracingLibrary::Builder& RayTracingLibrary::Builder::PushAnyHit(ShaderPtr<RHIShader> shader) noexcept
{
	assert(shader->GetShaderStage() == VK_SHADER_STAGE_ANY_HIT_BIT_KHR);
	assert(mImpl->tempCroup.anyHitShader == VK_SHADER_UNUSED_KHR);
	mImpl->tempCroup.anyHitShader = GetShaderIndex(shader);
	return *this;
}

RayTracingLibrary::Builder& RayTracingLibrary::Builder::PushIntersection(ShaderPtr<RHIShader> shader) noexcept
{
	assert(shader->GetShaderStage() == VK_SHADER_STAGE_INTERSECTION_BIT_KHR);
	assert(mImpl->tempCroup.intersectionShader == VK_SHADER_UNUSED_KHR);
	mImpl->tempCroup.intersectionShader = GetShaderIndex(shader);
	return *this;
}

RayTracingLibrary::Builder& RayTracingLibrary::Builder::PushCallable(ShaderPtr<RHIShader> shader) noexcept
{
	assert(shader->GetShaderStage() == VK_SHADER_STAGE_CALLABLE_BIT_KHR);
	assert(mImpl->tempCroup.generalShader == VK_SHADER_UNUSED_KHR);
	mImpl->tempCroup.generalShader = GetShaderIndex(shader);
	return *this;
}

void RayTracingLibrary::Builder::PushData(uint8 * data, uint64 dataSize)
{
	assert(mImpl->dataPtr == nullptr);
	mImpl->dataPtr = data;
	mImpl->dataSize = dataSize;
}

RayTracingLibrary::Builder& RayTracingLibrary::Builder::PackageGroup() noexcept
{
	RayTracingGroupType groupType = RayTracingGroupType::Raygen;

	if (mImpl->tempCroup.generalShader != VK_SHADER_UNUSED_KHR)
	{
		// Raygen, Miss, Callable only set at generalShader and the type are VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR
		mImpl->tempCroup.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR;
		// Find the shader group type
		switch (mImpl->shaderStages[mImpl->tempCroup.generalShader]->GetShaderStage())
		{
		case VK_SHADER_STAGE_RAYGEN_BIT_KHR:
			groupType = RayTracingGroupType::Raygen;
			break;
		case VK_SHADER_STAGE_MISS_BIT_KHR:
			groupType = RayTracingGroupType::Miss;
			break;
		case VK_SHADER_STAGE_CALLABLE_BIT_KHR:
			groupType = RayTracingGroupType::Callable;
			break;
		default:
			assert(0);
			break;
		}
	}
	else
	{
		if (mImpl->tempCroup.intersectionShader != VK_SHADER_UNUSED_KHR)
		{
			// if intersectionShader set, type must be VK_RAY_TRACING_SHADER_GROUP_TYPE_PROCEDURAL_HIT_GROUP_KHR
			mImpl->tempCroup.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_PROCEDURAL_HIT_GROUP_KHR;
		}
		else
		{
			// if closestHitShader and anyHitShader both set or only set one no matter witch, type always be VK_RAY_TRACING_SHADER_GROUP_TYPE_TRIANGLES_HIT_GROUP_KHR
			// closestHitShader and anyHitShader 两个都设置或者只设置其中一个, 都是Hit Group类型
			mImpl->tempCroup.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_TRIANGLES_HIT_GROUP_KHR;
		}
		groupType = RayTracingGroupType::Hit;
	}

	if (mImpl->dataPtr != nullptr)
	{
		mImpl->dataGroups[groupType].insert({ static_cast<uint32>(mImpl->shaderGroups.size()), std::vector<uint8>(mImpl->dataPtr, mImpl->dataPtr + mImpl->dataSize) });
	}

	mImpl->shaderGroups.push_back(mImpl->tempCroup);

	// reset tempGroup
	mImpl->tempCroup.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR;
	mImpl->tempCroup.generalShader = VK_SHADER_UNUSED_KHR;
	mImpl->tempCroup.closestHitShader = VK_SHADER_UNUSED_KHR;
	mImpl->tempCroup.anyHitShader = VK_SHADER_UNUSED_KHR;
	mImpl->tempCroup.intersectionShader = VK_SHADER_UNUSED_KHR;

	mImpl->dataPtr = nullptr;
	mImpl->dataSize = 0;

	return *this;
}

RayTracingLibrary::Builder& RayTracingLibrary::Builder::SetMaxPipelineRayRecursionDepth(uint32 maxPipelineRayRecursionDepth) noexcept
{
	mImpl->maxPipelineRayRecursionDepth = maxPipelineRayRecursionDepth;
	return *this;
}

RayTracingLibrary::Builder& RayTracingLibrary::Builder::SetMaxPipelineRayHitAttributeSize(uint32 maxPipelineRayHitAttributeSize) noexcept
{
	mImpl->maxPipelineRayHitAttributeSize = maxPipelineRayHitAttributeSize;
	return *this;
}

RayTracingLibrary::Builder& RayTracingLibrary::Builder::SetMaxPipelineRayPayloadSize(uint32 maxPipelineRayPayloadSize) noexcept
{
	mImpl->maxPipelineRayPayloadSize = maxPipelineRayPayloadSize;
	return *this;
}

RefCountPtr<RayTracingLibrary> RayTracingLibrary::Builder::Build()
{
	assert(IsInRenderThread());
	return RefCountPtr<RayTracingLibrary>(NewObject<RayTracingLibrary>(*mImpl));
}

RefCountPtr<RayTracingLibrary> RayTracingLibrary::Builder::Build(RenderProcess* process, const std::string& name)
{
	assert(IsInRenderThread());
	RayTracingLibrary* rayTracingLibrary = NewObject<RayTracingLibrary>(*mImpl);
	process->AddPipeline(name, rayTracingLibrary);
	return RefCountPtr<RayTracingLibrary>(rayTracingLibrary);
}

void RayTracingLibrary::ReleaseResource()
{
	if (IsInRenderThread())
	{
		if (mProcess != nullptr)
		{
			mProcess->RemovePipeline(mName, this);
		}
		DeleteObject<RayTracingLibrary>(this);
	}
	else
	{
		AsyncTask(GRenderThreadId, [this]() {
			if (mProcess != nullptr)
			{
				mProcess->RemovePipeline(mName, this);
			}
			DeleteObject<RayTracingLibrary>(this);
			});
	}
}

std::vector<VkRayTracingShaderGroupCreateInfoKHR>& RayTracingLibrary::GetShaderGroups()
{
	return mShaderGroups;
}

std::array<std::unordered_map<uint32, std::vector<uint8>>, 4>& RayTracingLibrary::GetDataGroups()
{
	return mDataGroups;
}

VkRayTracingPipelineInterfaceCreateInfoKHR RayTracingLibrary::GetRayTracingPipelineInterface() const
{
	return mRayTracingPipelineInterface;
}

uint32 RayTracingLibrary::GetMaxPipelineRayRecursionDepth() const
{
	return mMaxPipelineRayRecursionDepth;
}

std::vector<ShaderPtr<RHIShader>>& RayTracingLibrary::GetShaderStages()
{
	return mShaderStages;
}

uint32 RayTracingLibrary::Builder::GetShaderIndex(ShaderPtr<RHIShader> shader) const
{
	for (uint32 i = 0; i < mImpl->shaderStages.size(); ++i)
	{
		if (mImpl->shaderStages[i] == shader)
		{
			return i;
		}
	}
	// no shader exist
	// make sure the shader is add to shaderStages before make shaderGroup
	assert(0);
	return 0;
}

RayTracingLibrary::RayTracingLibrary(BuilderDetails & builderDetails)
{
	mDescriptor = builderDetails.descriptor;

	assert(builderDetails.maxPipelineRayRecursionDepth > 0 && builderDetails.maxPipelineRayRecursionDepth <= GVulkanRHI->GetRayTracingPipelineProperties().maxRayRecursionDepth);
	// maybe it can be zero, but i dont think so
	assert(builderDetails.maxPipelineRayHitAttributeSize != 0 && builderDetails.maxPipelineRayPayloadSize != 0);

	mShaderStages.assign(builderDetails.shaderStages.begin(), builderDetails.shaderStages.end());
	mShaderGroups.assign(builderDetails.shaderGroups.begin(), builderDetails.shaderGroups.end());
	for (uint32 i = 0; i < 4; ++i)
	{
		for (const auto& data : builderDetails.dataGroups[i])
		{
			mDataGroups[i].insert(data);
		}
	}

	mMaxPipelineRayRecursionDepth = builderDetails.maxPipelineRayRecursionDepth;

	mRayTracingPipelineInterface = {};
	mRayTracingPipelineInterface.sType = VK_STRUCTURE_TYPE_RAY_TRACING_PIPELINE_INTERFACE_CREATE_INFO_KHR;
	mRayTracingPipelineInterface.maxPipelineRayHitAttributeSize = builderDetails.maxPipelineRayHitAttributeSize;
	mRayTracingPipelineInterface.maxPipelineRayPayloadSize = builderDetails.maxPipelineRayPayloadSize;

	VkRayTracingPipelineCreateInfoKHR rayTracingPipelineCreateInfo{};
	rayTracingPipelineCreateInfo.sType = VK_STRUCTURE_TYPE_RAY_TRACING_PIPELINE_CREATE_INFO_KHR;
	// use this to create library
	rayTracingPipelineCreateInfo.flags = VK_PIPELINE_CREATE_LIBRARY_BIT_KHR;
	// target pipeline must use the same pipelineLayout
	// but it is no need to use the same ptr of pipelineLayout, you can create a new one with the same layout
	rayTracingPipelineCreateInfo.layout = mDescriptor->GetPipelineLayout();
	// target pipeline must use the same interface
	// but it is no need to use the same ptr of interface
	rayTracingPipelineCreateInfo.pLibraryInterface = &mRayTracingPipelineInterface;
	// As for the interface the maximum recursion depth must also be consistent across the pipeline
	rayTracingPipelineCreateInfo.maxPipelineRayRecursionDepth = mMaxPipelineRayRecursionDepth;

	std::vector<VkPipelineShaderStageCreateInfo> shadeStages;
	for (auto& shader : mShaderStages)
		shadeStages.push_back(shader->GetShaderStageCreateInfo());

	rayTracingPipelineCreateInfo.stageCount = static_cast<uint32>(shadeStages.size());
	rayTracingPipelineCreateInfo.pStages = shadeStages.data();
	rayTracingPipelineCreateInfo.groupCount = static_cast<uint32>(mShaderGroups.size());
	rayTracingPipelineCreateInfo.pGroups = mShaderGroups.data();

	VK_CHECK_RESULT(vkCreateRayTracingPipelinesKHR(GVulkanRHI->GetLogicalDevice(), VK_NULL_HANDLE, VK_NULL_HANDLE, 1, &rayTracingPipelineCreateInfo, nullptr, &mPipeline));
}

